# 模型菜单相关操作
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import os
from ruamel.yaml import YAML  # 使用ruamel.yaml替代yaml


def update_predict_yaml(shared_data):
    """更新predict.yaml中的配置参数（使用固定路径）"""
    # 固定路径（根据实际情况修改，例如"config/predict.yaml"）
    yaml_file = "../ultralytics/cfg/predict.yaml"

    # 初始化ruamel.yaml对象
    yaml = YAML()
    yaml.preserve_quotes = True

    try:
        data = {}
        if os.path.exists(yaml_file):
            with open(yaml_file, 'r', encoding='utf-8') as f:
                data = yaml.load(f) or {}
        else:
            data = {}

        # 更新配置（保留原有注释和其他未更改的配置）
        if shared_data.model_config.get("weight_path"):
            data['model'] = os.path.abspath(shared_data.model_config["weight_path"])
        data['conf'] = shared_data.model_config["confidence"]
        data['iou'] = shared_data.model_config["iou"]
        data['device'] = 0 if shared_data.model_config["device"] == "GPU" else 'cpu'
        data['imgsz'] = shared_data.model_config["image_size"]

        # 写回文件（保留注释）
        with open(yaml_file, 'w', encoding='utf-8') as f:
            yaml.dump(data, f)
        return True
    except Exception as e:
        print(f"更新配置文件失败: {str(e)}")
        return False


def update_scale_label(app, value, label, shared_data):
    """更新滑块对应的标签显示值"""
    label.config(text=f"{float(value):.1f}")
    if label == shared_data.ui_elements["confidence_label"]:
        shared_data.model_config["confidence"] = float(value)
        update_predict_yaml(shared_data)  # 同步更新YAML
    elif label == shared_data.ui_elements["iou_label"]:
        shared_data.model_config["iou"] = float(value)
        update_predict_yaml(shared_data)  # 同步更新YAML


def on_model_version_change(app, event, shared_data):
    """模型版本改变事件"""
    shared_data.model_config["version"] = app.model_version.get()
    from ui_utils import update_weight_combobox
    update_weight_combobox(app, shared_data)
    app.show_status(f"已选择模型版本: {shared_data.model_config['version']}")


def on_model_weight_change(app, event, shared_data):
    """模型权重改变事件"""
    selected_weight = app.model_weight.get()
    shared_data.model_config["weight_path"] = selected_weight
    app.show_status(f"已选择模型权重: {os.path.basename(selected_weight)}")
    update_predict_yaml(shared_data)  # 同步更新YAML


def cmd_load_model(shared_data, app):
    """加载自定义YOLO权重文件"""
    file_path = filedialog.askopenfilename(
        title="选择模型权重文件",
        filetypes=[
            ("YOLO模型文件", "*.pt *.onnx"),
            ("PyTorch模型", "*.pt"),
            ("ONNX模型", "*.onnx"),
            ("所有文件", "*.*")
        ]
    )
    if not file_path:
        return

    # 更新模型配置
    shared_data.model_config["weight_path"] = file_path

    # 更新下拉框显示
    current_version = app.model_version.get()
    shared_data.pretrained_models[current_version].insert(0, file_path)
    from ui_utils import update_weight_combobox
    update_weight_combobox(app, shared_data)
    shared_data.ui_elements["weight_combobox"].current(0)

    # 更新YAML并显示状态
    if update_predict_yaml(shared_data):
        app.show_status(f"已加载模型并更新配置: {os.path.basename(file_path)}")
    else:
        app.show_status(f"已加载模型，但配置更新失败: {os.path.basename(file_path)}")


def cmd_model_library(shared_data, app):
    """显示预训练模型列表"""
    dialog = tk.Toplevel(app)
    dialog.title("模型库")
    dialog.geometry("400x500")
    dialog.transient(app)
    dialog.grab_set()

    # 创建模型类别选择
    ttk.Label(dialog, text="选择模型版本:").pack(anchor=tk.W, padx=20, pady=(15, 5))
    model_var = tk.StringVar(value=shared_data.model_config["version"])
    model_combobox = ttk.Combobox(
        dialog,
        textvariable=model_var,
        values=list(shared_data.pretrained_models.keys()),
        state="readonly",
        width=30
    )
    model_combobox.pack(anchor=tk.W, padx=20, pady=(0, 10))

    # 创建模型列表
    frame = ttk.Frame(dialog)
    frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=10)

    listbox = tk.Listbox(frame, width=40, height=15)
    listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

    scrollbar = ttk.Scrollbar(frame, orient=tk.VERTICAL, command=listbox.yview)
    scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
    listbox.config(yscrollcommand=scrollbar.set)

    # 更新模型列表的函数
    def update_model_list(event=None):
        listbox.delete(0, tk.END)
        version = model_var.get()
        for model in shared_data.pretrained_models.get(version, []):
            listbox.insert(tk.END, model)

    # 初始加载模型列表
    update_model_list()
    model_combobox.bind("<<ComboboxSelected>>", update_model_list)

    # 下载/切换按钮
    def select_model():
        selection = listbox.curselection()
        if not selection:
            messagebox.showwarning("提示", "请选择一个模型")
            return

        index = selection[0]
        version = model_var.get()
        selected_model = shared_data.pretrained_models[version][index]

        # 更新主界面配置
        app.model_version.set(version)
        shared_data.model_config["version"] = version
        from ui_utils import update_weight_combobox
        update_weight_combobox(app, shared_data)
        shared_data.ui_elements["weight_combobox"].set(selected_model)
        shared_data.model_config["weight_path"] = selected_model

        # 更新YAML配置
        update_predict_yaml(shared_data)
        app.show_status(f"已切换到模型: {selected_model}")
        dialog.destroy()

    ttk.Button(dialog, text="下载并使用", command=select_model).pack(pady=20)

    dialog.wait_window()


def cmd_model_params(shared_data, app):
    """模型参数设置对话框"""
    dialog = tk.Toplevel(app)
    dialog.title("模型参数设置")
    dialog.geometry("400x350")
    dialog.transient(app)
    dialog.grab_set()

    # 置信度阈值
    ttk.Label(dialog, text="置信度阈值:").pack(anchor=tk.W, padx=20, pady=(15, 2))
    conf_var = tk.DoubleVar(value=shared_data.model_config["confidence"])
    conf_scale = ttk.Scale(
        dialog,
        from_=0.1,
        to=1.0,
        orient=tk.HORIZONTAL,
        length=300,
        variable=conf_var,
        command=lambda v: conf_label.config(text=f"{float(v):.1f}")
    )
    conf_scale.pack(anchor=tk.W, padx=20, pady=(0, 2))
    conf_label = ttk.Label(dialog, text=f"{shared_data.model_config['confidence']:.1f}")
    conf_label.pack(anchor=tk.W, padx=20, pady=(0, 10))

    # IOU阈值
    ttk.Label(dialog, text="IOU阈值:").pack(anchor=tk.W, padx=20, pady=(5, 2))
    iou_var = tk.DoubleVar(value=shared_data.model_config["iou"])
    iou_scale = ttk.Scale(
        dialog,
        from_=0.1,
        to=1.0,
        orient=tk.HORIZONTAL,
        length=300,
        variable=iou_var,
        command=lambda v: iou_label.config(text=f"{float(v):.1f}")
    )
    iou_scale.pack(anchor=tk.W, padx=20, pady=(0, 2))
    iou_label = ttk.Label(dialog, text=f"{shared_data.model_config['iou']:.1f}")
    iou_label.pack(anchor=tk.W, padx=20, pady=(0, 10))

    # 检测设备
    ttk.Label(dialog, text="检测设备:").pack(anchor=tk.W, padx=20, pady=(5, 2))
    device_var = tk.StringVar(value=shared_data.model_config["device"])
    devices = ["CPU"]
    if shared_data.model_config["device"] == "GPU":
        devices.append("GPU")
    device_combobox = ttk.Combobox(
        dialog,
        textvariable=device_var,
        values=devices,
        state="readonly",
        width=10
    )
    device_combobox.pack(anchor=tk.W, padx=20, pady=(0, 10))

    # 输入尺寸
    ttk.Label(dialog, text="输入尺寸:").pack(anchor=tk.W, padx=20, pady=(5, 2))
    size_var = tk.IntVar(value=shared_data.model_config["image_size"])
    size_combobox = ttk.Combobox(
        dialog,
        textvariable=size_var,
        values=[320, 480, 640, 800, 1024],
        state="readonly",
        width=10
    )
    size_combobox.pack(anchor=tk.W, padx=20, pady=(0, 10))

    # 确认按钮
    def confirm():
        # 更新配置
        shared_data.model_config["confidence"] = conf_var.get()
        shared_data.model_config["iou"] = iou_var.get()
        shared_data.model_config["device"] = device_var.get()
        shared_data.model_config["image_size"] = size_var.get()

        # 更新主界面滑块
        shared_data.ui_elements["confidence_scale"].set(conf_var.get())
        shared_data.ui_elements["confidence_label"].config(text=f"{conf_var.get():.1f}")
        shared_data.ui_elements["iou_scale"].set(iou_var.get())
        shared_data.ui_elements["iou_label"].config(text=f"{iou_var.get():.1f}")

        # 更新YAML配置
        if update_predict_yaml(shared_data):
            app.show_status("模型参数已更新并同步至配置文件")
        else:
            app.show_status("模型参数已更新，但配置文件同步失败")
        dialog.destroy()

    ttk.Button(dialog, text="确认", command=confirm).pack(pady=10)

    dialog.wait_window()
